import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import Delaunay

def generate_quality_mesh(size=100.0, h=5.0):
    """
    生成高质量三角形网格
    参数:
        size: 正方形边长的一半
        h: 目标网格尺寸
    """
    # 计算节点数
    n = int(2 * size / h) + 1
    
    # 生成结构化网格点
    x = np.linspace(-size, size, n)
    y = np.linspace(-size, size, n)
    X, Y = np.meshgrid(x, y)
    
    # 将网格点转换为节点列表
    nodes = np.column_stack((X.flatten(), Y.flatten()))
    
    # 添加内部扰动以改善三角形质量
    mask = (np.abs(nodes[:, 0]) < size-h) & (np.abs(nodes[:, 1]) < size-h)
    rand_scale = h * 0.1  # 控制扰动幅度
    nodes[mask] += np.random.rand(np.sum(mask), 2) * rand_scale
    
    # 使用Delaunay三角剖分
    tri = Delaunay(nodes)
    elements = tri.simplices
    
    # 计算并检查三角形质量
    quality = calculate_mesh_quality(nodes, elements)
    print(f"\n网格质量统计:")
    print(f"  最小质量: {np.min(quality):.3f}")
    print(f"  平均质量: {np.mean(quality):.3f}")
    print(f"  最大质量: {np.max(quality):.3f}")
    
    # 标记边界节点
    boundary_nodes = []
    tol = h/10
    
    for i, (x, y) in enumerate(nodes):
        if abs(x + size) < tol:  # 左边界
            boundary_nodes.append((i+1, 1))
        elif abs(x - size) < tol:  # 右边界
            boundary_nodes.append((i+1, 2))
        elif abs(y + size) < tol:  # 下边界
            boundary_nodes.append((i+1, 3))
        elif abs(y - size) < tol:  # 上边界
            boundary_nodes.append((i+1, 4))
    
    # 检查网格覆盖范围
    print("\n网格覆盖范围:")
    print(f"X: [{np.min(nodes[:,0]):.1f}, {np.max(nodes[:,0]):.1f}]")
    print(f"Y: [{np.min(nodes[:,1]):.1f}, {np.max(nodes[:,1]):.1f}]")
    
    # 检查边界节点分布
    print("\n边界节点分布:")
    for bc_type, name in [(1, "左"), (2, "右"), (3, "下"), (4, "上")]:
        nodes_of_type = [(i, x, y) for i, ((x, y), (idx, t)) in enumerate(zip(nodes, boundary_nodes)) if t == bc_type]
        if nodes_of_type:
            print(f"{name}边界节点范围:")
            print(f"  X: [{min(x for _,x,_ in nodes_of_type):.1f}, {max(x for _,x,_ in nodes_of_type):.1f}]")
            print(f"  Y: [{min(y for _,_,y in nodes_of_type):.1f}, {max(y for _,_,y in nodes_of_type):.1f}]")
            print(f"  节点数: {len(nodes_of_type)}")
    
    # 检查网格密度
    print("\n网格密度统计:")
    print(f"目标网格尺寸: {h:.1f}")
    print(f"节点总数: {len(nodes)}")
    print(f"单元总数: {len(elements)}")
    print(f"边界节点数: {len(boundary_nodes)}")
    
    return nodes, elements, boundary_nodes

def calculate_mesh_quality(nodes, elements):
    """
    计算网格质量（使用三角形形状因子）
    质量范围: [0,1], 1表示等边三角形
    """
    quality = []
    for element in elements:
        # 获取三角形顶点
        vertices = nodes[element]
        
        # 计算边长
        edges = np.array([
            vertices[1] - vertices[0],
            vertices[2] - vertices[1],
            vertices[0] - vertices[2]
        ])
        edge_lengths = np.sqrt(np.sum(edges**2, axis=1))
        
        # 计算面积
        a, b, c = edge_lengths
        s = (a + b + c) / 2  # 半周长
        area = np.sqrt(s*(s-a)*(s-b)*(s-c))
        
        # 计算质量因子 (4*sqrt(3)*area)/(sum(edge_lengths^2))
        quality_factor = 4 * np.sqrt(3) * area / np.sum(edge_lengths**2)
        quality.append(quality_factor)
    
    return np.array(quality)

def plot_mesh(nodes, elements):
    """绘制网格并显示质量"""
    plt.figure(figsize=(12, 5))
    
    # 绘制网格
    plt.subplot(121)
    plt.triplot(nodes[:, 0], nodes[:, 1], elements, 'b-', linewidth=0.5)
    plt.plot(nodes[:, 0], nodes[:, 1], 'r.', markersize=3)
    plt.axis('equal')
    plt.grid(True)
    plt.title('Triangle Mesh')
    
    # 绘制质量分布
    plt.subplot(122)
    quality = calculate_mesh_quality(nodes, elements)
    
    # 将单元质量映射到节点
    node_quality = np.zeros(len(nodes))
    node_count = np.zeros(len(nodes))
    
    for i, element in enumerate(elements):
        node_quality[element] += quality[i]
        node_count[element] += 1
    
    # 计算节点的平均质量
    node_quality = node_quality / np.maximum(node_count, 1)
    
    plt.tripcolor(nodes[:, 0], nodes[:, 1], elements, node_quality,
                 shading='gouraud', cmap='viridis')
    plt.colorbar(label='Element Quality')
    plt.axis('equal')
    plt.title('Mesh Quality Distribution')
    
    plt.tight_layout()
    plt.show()

def export_mesh(nodes, elements, boundary_nodes, filename_prefix='square_mesh'):
    """导出网格数据"""
    # 导出节点坐标
    with open(f'{filename_prefix}_nodes1.txt', 'w', encoding='utf-8-sig') as f:
        f.write(f'节点总数 {len(nodes)}\n')
        f.write('节点编号 X坐标 Y坐标\n')
        for i, (x, y) in enumerate(nodes):
            f.write(f'{i+1} {x:.6f} {y:.6f}\n')
    
    # 导出单元连接关系
    with open(f'{filename_prefix}_elements1.txt', 'w', encoding='utf-8-sig') as f:
        f.write(f'单元总数 {len(elements)}\n')
        f.write('单元编号 节点1 节点2 节点3\n')
        for i, element in enumerate(elements):
            f.write(f'{i+1} {element[0]+1} {element[1]+1} {element[2]+1}\n')
    
    # 导出边界节点信息
    with open(f'{filename_prefix}_boundary1.txt', 'w', encoding='utf-8-sig') as f:
        f.write(f'边界节点总数 {len(boundary_nodes)}\n')
        f.write('节点编号 边界类型(1:左边界 2:右边界 3:下边界 4:上边界)\n')
        for node_id, bc_type in boundary_nodes:
            f.write(f'{node_id} {bc_type}\n')

if __name__ == "__main__":
    size = 100.0
    h = 5
    nodes, elements, boundary_nodes = generate_quality_mesh(size=size, h=h)
    plot_mesh(nodes, elements)
    export_mesh(nodes, elements, boundary_nodes, 'square_mesh') 